#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import json
import subprocess
import argparse
import datetime
import shutil
import logging
import re
from typing import Dict, Any, List, Optional
from collections import defaultdict
import copy

# --- 全局配置 (保持不变) ---
LIB_DIR = "lib"
SETUP_SCRIPT = os.path.join(LIB_DIR, "setup_mysql_instances.sh")
SYSBENCH_SCRIPT = os.path.join(LIB_DIR, "run_sysbench_benchmark.sh")
TEMP_CONFIG_DIR = "temp_configs"

# --- 日志配置 (保持不变) ---
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    stream=sys.stdout)

# --- 辅助函数 ---

# [FIX v3.1] 修改 run_command 以兼容 Python 3.6
def run_command(command: List[str], log_file: str) -> bool:
    """执行一个shell命令，并将stdout/stderr实时输出到日志和控制台"""
    logging.info(f"执行命令: {' '.join(command)}")
    try:
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            # [MODIFIED] 使用 universal_newlines 替代 text
            universal_newlines=True,
            bufsize=1
        )
        with open(log_file, 'a') as f:
            for line in process.stdout:
                sys.stdout.write(line)
                f.write(line)
        
        process.wait()
        if process.returncode != 0:
            logging.error(f"命令执行失败，返回码: {process.returncode}")
            return False
        logging.info("命令执行成功。")
        return True
    except Exception as e:
        logging.error(f"执行命令时发生异常: {e}")
        return False

# [FIX v3.1] 修改 apply_os_params 和 revert_os_params 以兼容 Python 3.6
def apply_os_params(os_params: Dict[str, Any], devices: List[str]) -> Dict[str, str]:
    """应用操作系统参数，兼容旧版Python。"""
    original_values = {}
    logging.info("正在应用操作系统参数...")
    
    for key, value in os_params.items():
        try:
            if key.startswith("vm.") or key.startswith("net."):
                # [MODIFIED] 使用 subprocess.check_output 读取原始值
                original = subprocess.check_output(['sysctl', '-n', key]).decode('utf-8').strip()
                original_values[key] = original
                logging.info(f"  - 设置 {key} = {value} (原值: {original})")
                # [MODIFIED] 使用 subprocess.run 但不带 capture_output
                subprocess.run(['sudo', 'sysctl', '-w', f"{key}={value}"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

            elif key.startswith('/sys/'):
                for device_path in devices:
                    dev_name = os.path.basename(device_path)
                    actual_path = key.replace('nvme0n1', dev_name)
                    if os.path.exists(actual_path):
                        with open(actual_path, 'r') as f:
                            # [MODIFIED] 读取并解析原始值，去掉方括号
                            raw_original = f.read().strip()
                            match = re.search(r'\[(\w+)\]', raw_original)
                            original = match.group(1) if match else raw_original.split()[0]
                        if key not in original_values:
                           original_values[key] = original
                        logging.info(f"  - 设置 {actual_path} = {value} (原值: {original})")
                        with open(actual_path, 'w') as f:
                           f.write(str(value))
                    else:
                        logging.warning(f"  - 路径不存在，跳过: {actual_path}")
            else:
                 logging.warning(f"不支持的OS参数类型: {key}")
        except Exception as e:
            logging.error(f"设置OS参数 {key} 失败: {e}")
    return original_values

def revert_os_params(original_values: Dict[str, str], devices: List[str]):
    """恢复原始的操作系统参数，兼容旧版Python。"""
    if not original_values: return
    logging.info("正在恢复原始操作系统参数...")
    for key, value in original_values.items():
        try:
            if key.startswith("vm.") or key.startswith("net."):
                logging.info(f"  - 恢复 {key} = {value}")
                # [MODIFIED] 使用 subprocess.run 但不带 capture_output
                subprocess.run(['sudo', 'sysctl', '-w', f"{key}={value}"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            elif key.startswith('/sys/'):
                for device_path in devices:
                    dev_name = os.path.basename(device_path)
                    actual_path = key.replace('nvme0n1', dev_name)
                    if os.path.exists(actual_path):
                        logging.info(f"  - 恢复 {actual_path} = {value}")
                        with open(actual_path, 'w') as f:
                            f.write(str(value))
        except Exception as e:
            logging.error(f"恢复OS参数 {key} 失败: {e}")

# ... 此处是其他未修改的辅助函数 ...
def generate_mysql_config(base_configs: List[str], mysql_params: Dict[str, Any], output_path: str) -> bool:
    try:
        # [FIX v3.1] 确保目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'w') as f:
            f.write("# Auto-generated config for this test run\n\n")
            for base_path in base_configs:
                f.write(f"# Including base config: {base_path}\n")
                f.write(f"!include {os.path.abspath(base_path)}\n")
            f.write("\n[mysqld]\n")
            for key, value in mysql_params.items():
                f.write(f"{key:<35} = {value}\n")
        logging.info(f"已生成临时MySQL配置文件: {output_path}")
        return True
    except Exception as e:
        logging.error(f"生成MySQL配置文件失败: {e}")
        return False
        
def parse_sysbench_results(output_dir: str, num_instances: int, scenario: str, sysbench_args: str) -> float:
    total_tps = 0.0
    start_port = 3307
    threads_match = re.search(r'--threads=(\d+)', sysbench_args)
    time_match = re.search(r'--time=(\d+)', sysbench_args)
    threads = threads_match.group(1) if threads_match else 'default'
    time = time_match.group(1) if time_match else 'default'
    for i in range(num_instances):
        port = start_port + i
        log_filename = f"{scenario}_th{threads}_time{time}_result.log"
        log_path = os.path.join(output_dir, f"port_{port}", log_filename)
        if not os.path.exists(log_path): continue
        try:
            with open(log_path, 'r') as f:
                content = f.read()
                match = re.search(r"transactions:\s*\d+\s*\(([\d\.]+)\s*per sec\.\)", content)
                if match: total_tps += float(match.group(1))
        except Exception as e:
            logging.error(f"解析文件 {log_path} 失败: {e}")
    return total_tps

def generate_final_report(final_best_results: Dict, output_dir: str):
    report_path = os.path.join(output_dir, "final_analysis_report.txt")
    logging.info(f"正在生成最终报告: {report_path}")
    with open(report_path, 'w') as f:
        f.write("="*80 + "\n")
        f.write("=        MySQL Auto-Tuning Final Report (Control Variable Method)        =\n")
        f.write("="*80 + "\n\n")
        for workload, best_run in sorted(final_best_results.items()):
            f.write(f"--- Best Configuration for Workload: [{workload}] ---\n")
            f.write(f"  [+] Achieved Best Total TPS: {best_run['total_tps']:.2f}\n")
            f.write(f"  [+] Found after {best_run['total_runs']} test runs in {best_run['epoch']} epochs.\n\n")
            f.write("  Optimal Parameters Found:\n")
            f.write("  " + "-"*40 + "\n")
            params = best_run['params']
            f.write("    MySQL Parameters:\n")
            for k, v in params['mysql'].items():
                f.write(f"        {k:<30} = {v}\n")
            f.write("\n    OS Parameters:\n")
            for k, v in params['os'].items():
                f.write(f"        {k:<30} = {v}\n")
            f.write("\n    Sysbench Parameters:\n")
            for k, v in params['sysbench'].items():
                f.write(f"        {k:<30} = {v}\n")
            f.write("\n" + "="*80 + "\n\n")
    with open(report_path, 'r') as f:
        print("\n" + f.read())

# --- 核心函数 (run_single_test, main) 保持不变，因为问题出在辅助函数中 ---
def run_single_test(test_id: str, workload: str, params_to_test: Dict, config: Dict, run_log_file: str) -> float:
    mysql_inst_cfg = config['mysql_instances']
    devices = mysql_inst_cfg['devices'].split(',')
    num_instances = mysql_inst_cfg['count']
    
    logging.info(f"--- [Test ID: {test_id}] Running for workload '{workload}' ---")
    logging.info(f"Parameters: {json.dumps(params_to_test)}")

    original_os_params = {}
    try:
        original_os_params = apply_os_params(params_to_test['os'], devices)

        temp_config_path = os.path.join(TEMP_CONFIG_DIR, f"config_{test_id}.cnf")
        if not generate_mysql_config(config['mysql_base_configs'], params_to_test['mysql'], temp_config_path):
             raise Exception("Failed to generate MySQL config")

        setup_cmd = ['sudo', SETUP_SCRIPT, '-n', str(num_instances), '-c', mysql_inst_cfg['numa_nodes'], '-D', mysql_inst_cfg['devices'], '-f', os.path.abspath(temp_config_path)]
        if not run_command(setup_cmd, run_log_file):
            raise Exception("MySQL setup script failed")

        sysbench_dynamic_args = " ".join([f"--{k}={v}" for k, v in params_to_test['sysbench'].items()])
        full_sysbench_args = f"{config['sysbench_base_args']} {sysbench_dynamic_args}"
        ports_str = ",".join(map(str, range(3307, 3307 + num_instances)))

        logging.info("--> Phase: Sysbench prepare")
        prepare_cmd = [SYSBENCH_SCRIPT, '-p', ports_str, '-S', workload, '--prepare'] + full_sysbench_args.split()
        if not run_command(prepare_cmd, run_log_file):
            raise Exception("Sysbench prepare failed")

        logging.info(f"--> Phase: Sysbench run ({workload})")
        run_cmd = [SYSBENCH_SCRIPT, '-p', ports_str, '-S', workload] + full_sysbench_args.split()
        run_command(run_cmd, run_log_file) 

        total_tps = parse_sysbench_results('sysbench_results', num_instances, workload, full_sysbench_args)
        logging.info(f"RESULT for Test ID {test_id}: Workload[{workload}] | TPS: {total_tps:.2f}")

    except Exception as e:
        logging.error(f"Test run {test_id} failed: {e}")
        total_tps = 0.0
    finally:
        revert_os_params(original_os_params, devices)

    return total_tps
    
def main():
    parser = argparse.ArgumentParser(
        description="MySQL自动化性能优化器 (v3.1 - 兼容版)",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument('config_file', help="JSON格式的优化配置文件路径。")
    parser.add_argument('--dry-run', action='store_true', help="打印测试计划，不实际运行。")
    args = parser.parse_args()

    if not os.path.exists(args.config_file):
        logging.error(f"配置文件未找到: {args.config_file}")
        sys.exit(1)
    with open(args.config_file, 'r') as f:
        config = json.load(f)

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    top_level_dir = f"optimization_run_{timestamp}"
    os.makedirs(top_level_dir, exist_ok=True)
    os.makedirs(TEMP_CONFIG_DIR, exist_ok=True)
    shutil.copy(args.config_file, os.path.join(top_level_dir, "run_config.json"))
    run_log_file = os.path.join(top_level_dir, "full_run.log")
    file_handler = logging.FileHandler(run_log_file)
    file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    logging.getLogger().addHandler(file_handler)
    logging.info(f"所有输出将保存在: {top_level_dir}")

    params_to_optimize = []
    for group_name, params in config['optimization_parameters'].items():
        for param in params:
            if len(param['values']) > 1:
                params_to_optimize.append({'group': group_name, 'name': param['name'], 'values': param['values']})

    total_test_count_per_workload = 1 + sum(len(p['values']) - 1 for p in params_to_optimize)
    logging.info(f"优化策略: 控制变量法 (One-Variable-at-a-Time)")
    logging.info(f"待优化参数共 {len(params_to_optimize)} 个。")
    logging.info(f"每个 Workload 的测试总次数约为: {total_test_count_per_workload} * {config.get('optimization_epochs', 1)} epochs")

    if args.dry_run:
        print("--- [Dry Run] 测试计划 (控制变量法) ---")
        print(f"1. 运行基线测试。")
        for i, p in enumerate(params_to_optimize, 2):
            print(f"{i}. 优化参数 '{p['group']}.{p['name']}'，将测试 {len(p['values'])-1} 个新值。")
        sys.exit(0)

    final_best_results = {}
    
    for workload in config['workloads']:
        logging.info(f"============== 开始为 Workload: [{workload}] 进行优化 ==============")
        
        best_params_for_workload = defaultdict(dict)
        for group_name, params in config['optimization_parameters'].items():
            for param in params:
                best_params_for_workload[group_name][param['name']] = param['values'][0]
        
        best_tps_for_workload = 0.0
        total_runs_for_workload = 0

        for epoch in range(1, config.get('optimization_epochs', 1) + 1):
            logging.info(f"--- [Workload: {workload}, Epoch: {epoch}] ---")
            
            if epoch == 1:
                logging.info("--- Running initial baseline test ---")
                test_id = f"{workload}-E{epoch}-baseline"
                baseline_tps = run_single_test(test_id, workload, best_params_for_workload, config, run_log_file)
                best_tps_for_workload = baseline_tps
                total_runs_for_workload += 1
                logging.info(f"Baseline TPS for Epoch {epoch}: {best_tps_for_workload:.2f}")

            for param_info in params_to_optimize:
                group = param_info['group']
                name = param_info['name']
                values = param_info['values']
                
                logging.info(f"--- Optimizing parameter: [{group}.{name}] ---")
                
                current_best_value = best_params_for_workload[group][name]
                
                for value in values:
                    if value == current_best_value:
                        continue

                    params_to_test = copy.deepcopy(best_params_for_workload)
                    params_to_test[group][name] = value
                    
                    # [FIX v3.1] 修复test_id中的非法字符 '/'
                    safe_name = name.replace('/', '_')
                    test_id = f"{workload}-E{epoch}-{group}.{safe_name}-{value}"
                    
                    current_tps = run_single_test(test_id, workload, params_to_test, config, run_log_file)
                    total_runs_for_workload += 1

                    if current_tps > best_tps_for_workload:
                        logging.info(f"*** NEW BEST FOUND for [{workload}]! TPS: {best_tps_for_workload:.2f} -> {current_tps:.2f} ***")
                        logging.info(f"*** Parameter '{group}.{name}' updated to '{value}' ***")
                        best_tps_for_workload = current_tps
                        best_params_for_workload = copy.deepcopy(params_to_test)
                        current_best_value = value
            
            logging.info(f"--- Epoch {epoch} for workload [{workload}] finished. Best TPS so far: {best_tps_for_workload:.2f} ---")

        final_best_results[workload] = {
            'total_tps': best_tps_for_workload,
            'params': best_params_for_workload,
            'epoch': config.get('optimization_epochs', 1),
            'total_runs': total_runs_for_workload
        }

    if os.path.exists("sysbench_results"):
        shutil.move("sysbench_results", os.path.join(top_level_dir, "sysbench_results"))
    if os.path.exists(TEMP_CONFIG_DIR):
        shutil.move(TEMP_CONFIG_DIR, os.path.join(top_level_dir, "generated_configs"))
    
    generate_final_report(final_best_results, top_level_dir)
    logging.info(f"优化过程完成。所有数据已保存在: {top_level_dir}")

if __name__ == '__main__':
    if os.geteuid() != 0:
        print("错误: 此脚本需要root权限来修改系统参数和管理服务。请使用 'sudo python3 autotune.py ...' 运行。")
        sys.exit(1)
    main()